`timescale 1ns / 1ps

`define SLICE(BIT, IDX) (``BIT``) * (``IDX``) + (``BIT``) - 1 : (``BIT``) * (``IDX``)

module norm_actv
#(
    parameter N_IO      = 1,
    parameter N_BK      = 4,
    parameter N_RP      = 2,
    parameter N_CH      = 32,
    parameter BIT_I     = 16,
    parameter BIT_B     = 16,
    parameter BIT_M     = 16,
    parameter BIT_O     = 8,
    parameter RSHIFT    = 16,
    parameter PE_TYPE   = "DSP",    // "DSP" or "LUT"
    parameter ROM_FILE  = "file.mem",
    parameter ROM_TYPE  = "auto",
    parameter ROM_LATENCY = 1
)
(
    input   clk,
    input   rst,
    
    input   i_vld,
    output  i_rdy,
    input   [N_IO*BIT_I-1 : 0]      i_data,
    
    output  o_vld,
    input   o_rdy,
    output  [N_IO*BIT_O-1 : 0]      o_data
);

localparam DATAPATH_LATENCY = 5;
localparam W_BM  = BIT_M + BIT_B;
localparam W_IN  = N_IO * BIT_I;
localparam W_OUT = N_IO * BIT_O;
localparam W_ROM = N_IO * W_BM;
localparam D_ROM = N_CH / N_IO;

wire    w_rdy;
wire    w_vld;
wire    [W_ROM-1 : 0]   w_data;

wire    in_rdy;
wire    in_vld;
wire    enable;

assign i_rdy = in_rdy;

assign in_vld = i_vld & w_vld;
assign in_rdy = in_vld & enable;
assign enable = ~o_vld | o_rdy;

shift_reg_rst #(
    .N          (DATAPATH_LATENCY),
    .WIDTH      (1)
) inst_vaild_chain
(
    .clk        (clk),
    .rst        (rst),
    .clken      (enable),
    .i_data     (in_vld),
    .o_data     ( o_vld)
);

genvar i;
generate
    for (i = 0; i < N_IO; i = i + 1)
    begin
        if (PE_TYPE == "DSP") begin
            norm_actv_datapath_dsp
            #(
                .BIT_I      (BIT_I),
                .BIT_B      (BIT_B),
                .BIT_M      (BIT_M),
                .BIT_O      (BIT_O),
                .RSHIFT     (RSHIFT)
            ) inst_datapath
            (
                .clk        (clk),
                .clken      (enable),
                .norm_in    ( w_data[ `SLICE(W_BM , i) ] ),
                .data_in    ( i_data[ `SLICE(BIT_I, i) ] ),
                .data_out   ( o_data[ `SLICE(BIT_O, i) ] )
            ); 
        end else if (PE_TYPE == "LUT") begin
            norm_actv_datapath_lut
            #(
                .BIT_I      (BIT_I),
                .BIT_B      (BIT_B),
                .BIT_M      (BIT_M),
                .BIT_O      (BIT_O),
                .RSHIFT     (RSHIFT)
            ) inst_datapath
            (
                .clk        (clk),
                .clken      (enable),
                .norm_in    ( w_data[ `SLICE(W_BM , i) ] ),
                .data_in    ( i_data[ `SLICE(BIT_I, i) ] ),
                .data_out   ( o_data[ `SLICE(BIT_O, i) ] )
            ); 
        end
    end

    if (N_RP == 1) begin
        zq_rom_loader #(
            .WIDTH          (W_ROM),
            .DEPTH          (D_ROM),
            .ROM_FILE       (ROM_FILE),
            .ROM_TYPE       (ROM_TYPE),
            .ROM_LATENCY    (ROM_LATENCY)
        ) weight_loader
        (
            .clk        ( clk ),
            .rst        ( rst ),

            .o_vld      (  w_vld  ),
            .o_rdy      ( in_rdy  ),
            .o_data     (  w_data )
        );
    end
    else begin
        zq_rom_advloader #(
            .N_HEI          (N_RP),
            .N_BLK          (N_BK),
            .WIDTH          (W_ROM),
            .DEPTH          (D_ROM),
            .ROM_FILE       (ROM_FILE),
            .ROM_TYPE       (ROM_TYPE),
            .ROM_LATENCY    (ROM_LATENCY)
        ) weight_loader
        (
            .clk        ( clk ),
            .rst        ( rst ),

            .o_vld      (  w_vld  ),
            .o_rdy      ( in_rdy  ),
            .o_data     (  w_data )
        );
    end
endgenerate

endmodule
